use std::time::Duration;
use criterion::{criterion_group, criterion_main, Criterion, BenchmarkId};
use protoactor_rs::{
    actor::{Actor, ActorContext, ActorError, Props},
    system::ActorSystem,
    middleware::MetricsMiddleware,
    dispatcher::Dispatchers,
};
use async_trait::async_trait;
use tokio::runtime::Runtime;
use futures::stream::{self, StreamExt};

// 基准测试消息
#[derive(Debug)]
struct Ping(usize);

#[derive(Debug)]
struct Pong(usize);

// 基准测试 Actor
struct BenchActor {
    count: usize,
}

impl BenchActor {
    fn new() -> Self {
        Self { count: 0 }
    }
}

#[async_trait]
impl Actor for BenchActor {
    type Context = ActorContext;

    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn std::any::Any + Send>) -> Result<(), ActorError> {
        if let Ok(ping) = msg.downcast::<Ping>() {
            ctx.respond(Pong(ping.0));
        }
        Ok(())
    }
}

// 本地 PingPong 基准测试 - 单个 Actor
fn bench_local_ping_pong(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("local_ping_pong");
    
    for size in [1000, 10000, 100000].iter() {
        group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| {
            b.iter(|| {
                rt.block_on(async {
                    let system = ActorSystem::new().unwrap();
                    let props = Props::new(Box::new(|| Box::new(BenchActor::new())))
                        .with_dispatcher(Dispatchers::throughput(100));
                    let pid = system.spawn(props).await.unwrap();

                    stream::iter(0..size)
                        .map(Ping)
                        .for_each_concurrent(None, |msg| async {
                            let _: Pong = system.request(&pid, msg, Duration::from_secs(1))
                                .await
                                .unwrap();
                        })
                        .await;
                });
            });
        });
    }
    group.finish();
}

// 本地 PingPong 基准测试 - 多个 Actor
fn bench_local_ping_pong_multi(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("local_ping_pong_multi");
    
    for actors in [10, 100, 1000].iter() {
        group.bench_with_input(BenchmarkId::from_parameter(actors), actors, |b, &actors| {
            b.iter(|| {
                rt.block_on(async {
                    let system = ActorSystem::new().unwrap();
                    let mut pids = Vec::with_capacity(actors);

                    // 创建多个actor
                    for _ in 0..actors {
                        let props = Props::new(Box::new(|| Box::new(BenchActor::new())))
                            .with_dispatcher(Dispatchers::throughput(100));
                        let pid = system.spawn(props).await.unwrap();
                        pids.push(pid);
                    }

                    // 并发发送消息
                    stream::iter(&pids)
                        .for_each_concurrent(None, |pid| async {
                            let _: Pong = system.request(pid, Ping(1), Duration::from_secs(1))
                                .await
                                .unwrap();
                        })
                        .await;
                });
            });
        });
    }
    group.finish();
}

// 远程 PingPong 基准测试
fn bench_remote_ping_pong(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("remote_ping_pong");
    
    // 减小测试规模以提高稳定性
    for size in [50, 500, 5000].iter() {
        group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| {
            // 预热系统
            rt.block_on(async {
                let system1 = ActorSystem::new().unwrap();
                let system2 = ActorSystem::new().unwrap();
                system1.start_remote("127.0.0.1:8001").await.unwrap();
                system2.start_remote("127.0.0.1:8002").await.unwrap();
                tokio::time::sleep(Duration::from_millis(100)).await;
            });

            b.iter(|| {
                rt.block_on(async {
                    let system1 = ActorSystem::new().unwrap();
                    let system2 = ActorSystem::new().unwrap();

                    system1.start_remote("127.0.0.1:8001").await.unwrap();
                    system2.start_remote("127.0.0.1:8002").await.unwrap();

                    let props = Props::new(Box::new(|| Box::new(BenchActor::new())))
                        .with_dispatcher(Dispatchers::throughput(100))
                        .with_middleware(MetricsMiddleware::new());
                    let pid = system2.spawn(props).await.unwrap();

                    stream::iter(0..size)
                        .map(Ping)
                        .for_each_concurrent(Some(32), |msg| async {  // 限制并发数
                            let _: Pong = system1.request(&pid, msg, Duration::from_secs(1))
                                .await
                                .unwrap();
                        })
                        .await;
                });
            });
        });
    }
    group.finish();
}

// SkyNet 基准测试
fn bench_skynet(c: &mut Criterion) {
    let rt = Runtime::new().unwrap();
    let mut group = c.benchmark_group("skynet");
    
    // 调整测试规模
    for size in [5000, 50000, 500000].iter() {
        group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| {
            b.iter(|| {
                rt.block_on(async {
                    let system = ActorSystem::new().unwrap();
                    let mut pids = Vec::with_capacity(size);

                    // 分批创建actor以减少内存压力
                    for chunk in (0..size).step_by(1000) {
                        let chunk_size = std::cmp::min(1000, size - chunk);
                        stream::iter(0..chunk_size)
                            .for_each_concurrent(Some(32), |_| async {
                                let props = Props::new(Box::new(|| Box::new(BenchActor::new())))
                                    .with_dispatcher(Dispatchers::throughput(100));
                                let pid = system.spawn(props).await.unwrap();
                                pids.push(pid);
                            })
                            .await;
                    }

                    // 分批发送消息
                    for chunk in pids.chunks(1000) {
                        stream::iter(chunk)
                            .for_each_concurrent(Some(32), |pid| async {
                                system.send(pid, Ping(1)).await.unwrap();
                            })
                            .await;
                    }
                });
            });
        });
    }
    group.finish();
}

criterion_group!(
    name = benches;
    config = Criterion::default()
        .sample_size(10)
        .measurement_time(Duration::from_secs(30))
        .warm_up_time(Duration::from_secs(5));  // 添加预热时间
    targets = bench_local_ping_pong, bench_local_ping_pong_multi, 
             bench_remote_ping_pong, bench_skynet
);
criterion_main!(benches); 